package adkr

import (
	"ebbflow/pkg/utils"
	"bytes"
	"encoding/json"
	"log"
)

// ReliableBroadcast is a generic implementation of Bracha's Reliable Broadcast protocol.
func ReliableBroadcast(
	sid string,
	pid, n, f, leader int,
	inputCh <-chan []byte,
	outputCh chan<- []byte,
	msgInCh <-chan Message,
	msgOutCh chan<- Message,
) {
	log.Printf("RBC-SID:%s PID:%d Starting for leader %d", sid, pid, leader)

	k := n - 2*f
	echoThreshold := n - f
	readyThreshold := f + 1
	outputThreshold := 2*f + 1

	sendTo := func(to int, subMsg RBCMessage) {
		payload, _ := json.Marshal(subMsg)
		msgOutCh <- Message{Tag: ADKR_MVBA, Round: leader, To: to, Payload: payload}
	}
	broadcast := func(subMsg RBCMessage) {
		payload, _ := json.Marshal(subMsg)
		msgOutCh <- Message{Tag: ADKR_MVBA, Round: leader, To: -1, Payload: payload}
	}

	if pid == leader {
		go func() {
            val := <-inputCh
            totalLen := len(val)
            stripes, err := utils.Encode(k, n, val)
			if err != nil {
				log.Printf("RBC-SID:%s PID:%d [Leader] FATAL: %v", sid, pid, err)
				return
			}
			tree, rootHash := utils.MerkleTree(stripes)

			for i := 0; i < n; i++ {
				branch, _ := utils.GetMerkleBranch(i, tree)
                valPayload := RBCValPayload{RootHash: rootHash, Branch: branch, Stripe: stripes[i], TotalLen: totalLen}
				payloadBytes, _ := json.Marshal(valPayload)
				sendTo(i, RBCMessage{Tag: RBCValTag, Payload: payloadBytes})
			}
			log.Printf("RBC-SID:%s PID:%d [Leader] Sent VAL messages with root hash %x", sid, pid, rootHash[:4])
		}()
	}

    var fromLeaderRootHash []byte
    stripes := make(map[string][][]byte)
    echoCounters := make(map[string]int)
    readyCounters := make(map[string]map[int]struct{})
    totalLens := make(map[string]int)
    readySent := false

	for msg := range msgInCh {
		var subMsg RBCMessage
		if err := json.Unmarshal(msg.Payload, &subMsg); err != nil {
			continue
		}
		sender := msg.Sender
		var rootHashForCheck []byte

		switch subMsg.Tag {
		case RBCValTag:
			if fromLeaderRootHash != nil {
				continue
			}
            var valPayload RBCValPayload
			json.Unmarshal(subMsg.Payload, &valPayload)
			if sender != leader {
				continue
			}
			if !utils.MerkleVerify(n, valPayload.Stripe, valPayload.RootHash, valPayload.Branch, pid) {
				continue
			}
            fromLeaderRootHash = valPayload.RootHash
            totalLens[string(valPayload.RootHash)] = valPayload.TotalLen
			echoPayloadBytes, _ := json.Marshal(valPayload)
			broadcast(RBCMessage{Tag: RBCEchoTag, Payload: echoPayloadBytes})

		case RBCEchoTag:
            var echoPayload RBCValPayload
			json.Unmarshal(subMsg.Payload, &echoPayload)
			rootHashHex := string(echoPayload.RootHash)
			rootHashForCheck = echoPayload.RootHash
			if !utils.MerkleVerify(n, echoPayload.Stripe, echoPayload.RootHash, echoPayload.Branch, sender) {
				continue
			}
			if _, ok := stripes[rootHashHex]; !ok {
				stripes[rootHashHex] = make([][]byte, n)
				readyCounters[rootHashHex] = make(map[int]struct{})
			}
			if stripes[rootHashHex][sender] != nil {
				continue
			}
            stripes[rootHashHex][sender] = echoPayload.Stripe
            if _, ok := totalLens[rootHashHex]; !ok && echoPayload.TotalLen > 0 {
                totalLens[rootHashHex] = echoPayload.TotalLen
            }
			echoCounters[rootHashHex]++
			if echoCounters[rootHashHex] >= echoThreshold && !readySent {
				readySent = true
				readyPayload := RBCReadyPayload{RootHash: echoPayload.RootHash}
				readyPayloadBytes, _ := json.Marshal(readyPayload)
				broadcast(RBCMessage{Tag: RBCReadyTag, Payload: readyPayloadBytes})
			}

		case RBCReadyTag:
			var readyPayload RBCReadyPayload
			json.Unmarshal(subMsg.Payload, &readyPayload)
			rootHashHex := string(readyPayload.RootHash)
			rootHashForCheck = readyPayload.RootHash
			if _, ok := readyCounters[rootHashHex]; !ok {
				readyCounters[rootHashHex] = make(map[int]struct{})
			}
			if _, ok := readyCounters[rootHashHex][sender]; ok {
				continue
			}
			readyCounters[rootHashHex][sender] = struct{}{}
			if len(readyCounters[rootHashHex]) >= readyThreshold && !readySent {
				readySent = true
				broadcast(RBCMessage{Tag: RBCReadyTag, Payload: subMsg.Payload})
			}
		}

		if rootHashForCheck != nil {
			rootHashHex := string(rootHashForCheck)
			if _, ok := readyCounters[rootHashHex]; !ok {
				continue
			}
			if len(readyCounters[rootHashHex]) >= outputThreshold {
				stripeCount := 0
				if _, ok := stripes[rootHashHex]; ok {
					for _, s := range stripes[rootHashHex] {
						if s != nil {
							stripeCount++
						}
					}
				}
				if stripeCount >= k {
                    decoded, err := utils.Decode(k, n, stripes[rootHashHex])
					if err != nil {
						return
					}
                    // Trim to leader-declared total length if present
                    if L, ok := totalLens[rootHashHex]; ok && L > 0 && L <= len(decoded) {
                        decoded = decoded[:L]
                    }
					reEncoded, err := utils.Encode(k, n, decoded)
					if err != nil {
						return
					}
					_, reComputedRoot := utils.MerkleTree(reEncoded)
					if !bytes.Equal(reComputedRoot, rootHashForCheck) {
						return
					}
                    // Attempt truncation: parse any RBCValPayload we have to read TotalLen if included
                    var tmp RBCValPayload
                    if err := json.Unmarshal(subMsg.Payload, &tmp); err == nil && tmp.TotalLen > 0 && tmp.TotalLen <= len(decoded) {
                        decoded = decoded[:tmp.TotalLen]
                    }
                    log.Printf("RBC-SID:%s PID:%d Successfully decoded content from leader %d", sid, pid, leader)
                    outputCh <- decoded
					return
				}
			}
		}
	}
}
